import requests
from selenium import webdriver
import re
import utils

class Scan:
    
    def __init__(self, session, url='https://mp.weixin.qq.com'):
        self.url = url
        self.session = session
        self.options = webdriver.ChromeOptions()
        # self.options.add_argument(r"user-data-dir=C:/Users/A/AppData/Local/Google/Chrome/User Data")
        self.driver = webdriver.Chrome('C:\chromedriver\chromedriver.exe', chrome_options=self.options)

    def get_session_with_cookie(self, gh_id, pwd):
        data = utils.get_post_data(gh_id, pwd)

        res = self.session.post(self.url + '/cgi-bin/bizlogin?action=startlogin', data=data, headers=utils.headers
                                , verify=False)
        c = requests.utils.dict_from_cookiejar(res.cookies) 

        self.driver.maximize_window()
        self.driver.get(self.url + res.json()['redirect_url'])

        for key in c :
            self.driver.add_cookie({'name': key, 'value': c[key]})
            self.session.cookies.set(key, c[key])

        self.driver.get('https://mp.weixin.qq.com' + res.json()['redirect_url'])

        token = []

        while 1:
            current_url = self.driver.current_url
            if re.search('token', current_url):
                token.append(re.search('token=(.*)', current_url).group(1))
                break

        for cookie in self.driver.get_cookies():
            self.session.cookies.set(cookie['name'], cookie['value'])

        self.driver.quit()

        return token[0]

if __name__ == '__main__':
    session = utils.requests_retry_session()
    scan = Scan(session)
    scan.get_session_with_cookie('gh_5367ec89faf8', 'qwerasdf66')